package com.example.screenprojection;

import android.hardware.display.DisplayManager;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.media.projection.MediaProjection;
import android.view.Surface;

import java.io.IOException;
import java.nio.ByteBuffer;

/**
 * 采用H.265编码 H.265/HEVC的编码架构大致上和H.264/AVC的架构相似 H.265又称为HEVC(全称High Efficiency Video Coding，高效率视频编码)
 */
public class ScreenEncoder extends Thread {

  private static final int VIDEO_WIDTH = 2160;
  private static final int VIDEO_HEIGHT = 3840;
  private static final int SCREEN_FRAME_RATE = 20;
  private static final int SCREEN_FRAME_INTERVAL = 1;
  private static final long SOCKET_TIME_OUT = 10000;
  // I帧
  private static final int TYPE_FRAME_INTERVAL = 19;
  // vps帧
  private static final int TYPE_FRAME_VPS = 32;
  private final MediaProjection mMediaProjection;
  private final SocketManager mSocketManager;
  private MediaCodec mMediaCodec;
  private boolean mPlaying = true;
  // 记录vps pps sps
  private byte[] vps_pps_sps;

  public ScreenEncoder (SocketManager socketManager, MediaProjection mediaProjection) {
    mSocketManager = socketManager;
    mMediaProjection = mediaProjection;
  }

  public void startEncode() {
    // 声明MediaFormat，创建视频格式。
    MediaFormat mediaFormat =
        MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_HEVC, VIDEO_WIDTH, VIDEO_HEIGHT);
    // 描述视频格式的内容的颜色格式
    mediaFormat.setInteger(
        MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatSurface);
    // 比特率（比特/秒）
    mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, VIDEO_WIDTH * VIDEO_HEIGHT);
    // 帧率
    mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, SCREEN_FRAME_RATE);
    // I帧的频率
    mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, SCREEN_FRAME_INTERVAL);
    try {
      // 创建编码MediaCodec 类型是video/hevc
      mMediaCodec = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_HEVC);
      // 配置编码器
      mMediaCodec.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
      // 创建一个目的surface来存放输入数据
      Surface surface = mMediaCodec.createInputSurface();
      // 获取屏幕流
      mMediaProjection.createVirtualDisplay(
          "screen",
          VIDEO_WIDTH,
          VIDEO_HEIGHT,
          1,
          DisplayManager.VIRTUAL_DISPLAY_FLAG_PUBLIC,
          surface,
          null,
          null);
    } catch (IOException e) {
      e.printStackTrace();
    }

    start();
  }

  @Override
  public void run() {
    // 编解码器立即进入刷新子状态
    mMediaCodec.start();
    // 缓存区的元数据
    MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
    // 子线程需要一直运行，进行编码推流，所以要一直循环
    while (mPlaying) {
      // 查询编码输出
      int outPutBufferId = mMediaCodec.dequeueOutputBuffer(bufferInfo, SOCKET_TIME_OUT);
      if (outPutBufferId >= 0) {
        // 获取编码之后的数据输出流队列
        ByteBuffer byteBuffer = mMediaCodec.getOutputBuffer(outPutBufferId);
        // 添加上vps,sps,pps
        encodeData(byteBuffer, bufferInfo);
        // 处理完成，释放ByteBuffer数据
        mMediaCodec.releaseOutputBuffer(outPutBufferId, false);
      }
    }
  }

  private void encodeData(ByteBuffer byteBuffer, MediaCodec.BufferInfo bufferInfo) {
    // 偏移4 00 00 00 01为分隔符需要跳过
    int offSet = 4;
    if (byteBuffer.get(2) == 0x01) {
      offSet = 3;
    }
    // 计算出当前帧的类型
    int type = (byteBuffer.get(offSet) & 0x7E) >> 1;
    if (type == TYPE_FRAME_VPS) {
      // 保存vps sps pps信息
      vps_pps_sps = new byte[bufferInfo.size];
      byteBuffer.get(vps_pps_sps);
    } else if (type == TYPE_FRAME_INTERVAL) {
      // 将保存的vps sps pps添加到I帧前
      final byte[] bytes = new byte[bufferInfo.size];
      byteBuffer.get(bytes);
      byte[] newBytes = new byte[vps_pps_sps.length + bytes.length];
      System.arraycopy(vps_pps_sps, 0, newBytes, 0, vps_pps_sps.length);
      System.arraycopy(bytes, 0, newBytes, vps_pps_sps.length, bytes.length);
      // 将重新编码好的数据发送出去
      mSocketManager.sendData(newBytes);
    } else {
      // B帧 P帧 直接发送
      byte[] bytes = new byte[bufferInfo.size];
      byteBuffer.get(bytes);
      mSocketManager.sendData(bytes);
    }
  }

  public void stopEncode() {
    mPlaying = false;
    if (mMediaCodec != null) {
      mMediaCodec.release();
    }
    if (mMediaProjection != null) {
      mMediaProjection.stop();
    }
  }
}
